import argparse
import os
import sys
from pathlib import Path
from collections import OrderedDict

import onnx
import numpy as np
import onnxruntime as ort

parser = argparse.ArgumentParser(
    description="dump outputs of all onnx nodes.")
parser.add_argument('-m', '--model_path', type=str, help='path to onnx model')
parser.add_argument('-d', '--data_npys', type=str, nargs='+',
                    help='path to input data, only receive npy file.')
parser.add_argument('-o', '--output_dir', type=str, 
                    help='path to save dumped data.')
args = parser.parse_args()

output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
model = onnx.load(model_path)

print("Adding all node outputs to model outputs...")
for node in model.graph.node:
    for i, output in enumerate(node.output):
        model.graph.output.extend([onnx.ValueInfoProto(name=output)])

print("Infering...")
sess = ort.InferenceSession(model.SerializeToString())
onnx_inputs = sess.get_inputs()
input_feed = {
    inp.name: np.load(args.data_npys[i]) 
    for i, inp in enumerate(onnx_inputs)
}
output_names = [x.name for x in sess.get_outputs()]
outputs = sess.run(output_names, input_feed)
out_dict = OrderedDict(zip(output_names, outputs))

print(f"Saving outputs data to {str(output_dir)}...")
for name, output in out_dict.items():
    name = name.replace('/', '+')
    np.save(str(output_dir/f'{name}.npy'), output)
